import cv2
import numpy as np

if __name__ == '__main__':
    # 1. 图片输入
    image_np = cv2.imread('qu_noise.jpg')
    if image_np is None:
        print("无法读取图片，请检查路径")
        exit()

    # 2. HSV空间转换
    hsv_image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2HSV)

    # 3. 分别创建红、绿、蓝、黄的HSV掩膜
    # 红色掩膜
    red_low1 = np.array([0, 43, 46])
    red_high1 = np.array([10, 255, 255])
    red_low2 = np.array([156, 43, 46])
    red_high2 = np.array([180, 255, 255])
    mask_red1 = cv2.inRange(hsv_image_np, red_low1, red_high1)
    mask_red2 = cv2.inRange(hsv_image_np, red_low2, red_high2)
    mask_red = cv2.bitwise_or(mask_red1, mask_red2)

    # 绿色掩膜
    green_low = np.array([35, 43, 46])
    green_high = np.array([77, 255, 255])
    mask_green = cv2.inRange(hsv_image_np, green_low, green_high)

    # 蓝色掩膜
    blue_low = np.array([100, 43, 46])
    blue_high = np.array([124, 255, 255])
    mask_blue = cv2.inRange(hsv_image_np, blue_low, blue_high)

    # 黄色掩膜
    yellow_low = np.array([26, 43, 46])
    yellow_high = np.array([34, 255, 255])
    mask_yellow = cv2.inRange(hsv_image_np, yellow_low, yellow_high)

    # 合并所有颜色掩膜
    mask_all = cv2.bitwise_or(mask_red, mask_green)
    mask_all = cv2.bitwise_or(mask_all, mask_blue)
    mask_all = cv2.bitwise_or(mask_all, mask_yellow)

    # 4. 形态学开运算（先腐蚀再膨胀，去除小噪点）
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    open_mask = cv2.morphologyEx(mask_all, cv2.MORPH_OPEN, kernel, iterations=3)

    # 5. 颜色替换（保留原始颜色，仅通过掩膜降噪）
    result = image_np.copy()
    for i in range(result.shape[0]):
        for j in range(result.shape[1]):
            if open_mask[i, j] == 255:
                result[i, j] = image_np[i, j]  # 保留原始颜色
            else:
                result[i, j] = (255, 255, 255)  # 噪点区域设为白色

    # 6. 图片输出
    cv2.imshow('result', result)
    cv2.imwrite('Denoising.png', result)
    cv2.waitKey(0)
